import os
import uuid
import time
import struct
import threading
import re
import logging
import platform
import gradio as gr
import shutil
import socket
from cryptography.fernet import Fernet, InvalidToken
from werkzeug.utils import secure_filename

# 配置日志 - 同时输出到控制台和文件
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

# 创建控制台处理器
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.INFO)

# 创建文件处理器
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
log_file = os.path.join(BASE_DIR, 'en_decryption.log')
file_handler = logging.FileHandler(log_file, mode='a')
file_handler.setLevel(logging.INFO)

# 创建日志格式
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
console_handler.setFormatter(formatter)
file_handler.setFormatter(formatter)

# 添加处理器到日志器
logger.addHandler(console_handler)
logger.addHandler(file_handler)

# 创建必要的目录
UPLOAD_FOLDER = os.path.join(BASE_DIR, 'user-upload')
ENCRYPTED_FOLDER = os.path.join(BASE_DIR, 'encrypted-file')
DECRYPTED_FOLDER = os.path.join(BASE_DIR, 'decrypted-file')

os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(ENCRYPTED_FOLDER, exist_ok=True)
os.makedirs(DECRYPTED_FOLDER, exist_ok=True)

# 安全文件名验证
def is_valid_filename(filename):
    return re.match(r'^[\w,\s-]+\.[A-Za-z]{2,4}$', filename) is not None

# 清理旧文件函数
def cleanup_old_files():
    try:
        now = time.time()
        for folder in [UPLOAD_FOLDER, ENCRYPTED_FOLDER, DECRYPTED_FOLDER]:
            for filename in os.listdir(folder):
                file_path = os.path.join(folder, filename)
                if os.path.isfile(file_path):
                    file_age = now - os.path.getmtime(file_path)
                    if file_age > 60:  # 60秒后删除
                        try:
                            os.remove(file_path)
                            logger.info(f"清理文件: {file_path}")
                        except:
                            logger.error(f"清理文件失败: {file_path}")
    except Exception as e:
        logger.error(f"清理文件时出错: {e}")

# 定时清理任务
def scheduled_cleanup():
    logger.info("清理线程已启动")
    while True:
        try:
            cleanup_old_files()
        except Exception as e:
            logger.error(f"清理过程中出错: {e}")
        time.sleep(30)

# 启动清理线程
cleanup_thread = threading.Thread(target=scheduled_cleanup, daemon=True)
cleanup_thread.start()
logger.info("清理线程已创建")

# 确保文件完全写入磁盘
def ensure_file_written(file_path):
    retry_count = 0
    while retry_count < 5:
        if os.path.exists(file_path) and os.path.getsize(file_path) > 0:
            return True
        time.sleep(0.1)
        retry_count += 1
    return False

# 加密文件处理函数
def encrypt_file(file_path):
    if file_path is None:
        return None, "未选择文件"

    filename = os.path.basename(file_path)
    logger.info(f"开始加密文件: {filename}")

    if not is_valid_filename(filename):
        return None, "文件名包含非法字符"

    if not filename.endswith('.md'):
        return None, "请上传 .md 文件"

    secure_file_name = secure_filename(filename)
    base_name = os.path.splitext(secure_file_name)[0]

    unique_id = uuid.uuid4().hex
    upload_path = os.path.join(UPLOAD_FOLDER, f"{unique_id}.md")
    encrypted_path = os.path.join(ENCRYPTED_FOLDER, f"{unique_id}.bin")

    output_filename = f"{base_name}_encrypted.bin"
    final_output_path = os.path.join(ENCRYPTED_FOLDER, output_filename)

    try:
        shutil.copy(file_path, upload_path)
        file_size = os.path.getsize(upload_path)
        logger.info(f"保存上传文件: {upload_path} (大小: {file_size}字节)")

        if file_size > 4 * 1024 * 1024:
            return None, "文件太大，最大支持4MB"

        key = Fernet.generate_key()
        fernet = Fernet(key)

        with open(upload_path, 'rb') as infile:
            file_data = infile.read()

        encrypted_data = fernet.encrypt(file_data)

        key_length = len(key)
        composite_data = struct.pack('>I', key_length) + key + encrypted_data

        with open(encrypted_path, 'wb') as outfile:
            outfile.write(composite_data)

        if not ensure_file_written(encrypted_path):
            logger.error(f"加密文件写入失败: {encrypted_path}")
            return None, "加密文件创建失败，请重试"

        os.rename(encrypted_path, final_output_path)

        encrypted_size = os.path.getsize(final_output_path)
        logger.info(f"保存加密文件: {final_output_path} (大小: {encrypted_size}字节)")

        return final_output_path, f"加密成功！文件名: {output_filename}"

    except Exception as e:
        logger.error(f"加密文件时出错: {str(e)}", exc_info=True)
        return None, f"加密文件时出错: {str(e)}"
    finally:
        try:
            if os.path.exists(upload_path):
                os.remove(upload_path)
                logger.info(f"删除上传文件: {upload_path}")
        except Exception as e:
            logger.error(f"删除上传文件失败: {str(e)}")

# 解密文件处理函数
def decrypt_file(file_path):
    if file_path is None:
        return None, "未选择文件"

    filename = os.path.basename(file_path)
    logger.info(f"开始解密文件: {filename}")

    if not is_valid_filename(filename):
        return None, "文件名包含非法字符"

    if not filename.endswith('.bin'):
        return None, "请上传 .bin 文件"

    secure_file_name = secure_filename(filename)
    base_name = os.path.splitext(secure_file_name)[0]

    original_name = base_name.replace("_encrypted", "")

    unique_id = uuid.uuid4().hex
    upload_path = os.path.join(UPLOAD_FOLDER, f"{unique_id}.bin")
    decrypted_path = os.path.join(DECRYPTED_FOLDER, f"{unique_id}.md")

    output_filename = f"{original_name}_decrypted.md"
    final_output_path = os.path.join(DECRYPTED_FOLDER, output_filename)

    try:
        shutil.copy(file_path, upload_path)
        file_size = os.path.getsize(upload_path)
        logger.info(f"保存上传文件: {upload_path} (大小: {file_size}字节)")

        if file_size < 4:
            return None, "无效的加密文件格式"

        with open(upload_path, 'rb') as infile:
            composite_data = infile.read()

        if len(composite_data) < 4:
            return None, "无效的加密文件格式"

        key_length = struct.unpack('>I', composite_data[:4])[0]

        if len(composite_data) < 4 + key_length:
            return None, "无效的加密文件格式"

        key = composite_data[4:4+key_length]
        encrypted_data = composite_data[4+key_length:]

        fernet = Fernet(key)
        decrypted_data = fernet.decrypt(encrypted_data)

        with open(decrypted_path, 'wb') as outfile:
            outfile.write(decrypted_data)

        if not ensure_file_written(decrypted_path):
            logger.error(f"解密文件写入失败: {decrypted_path}")
            return None, "解密文件创建失败，请重试"

        os.rename(decrypted_path, final_output_path)

        decrypted_size = os.path.getsize(final_output_path)
        logger.info(f"保存解密文件: {final_output_path} (大小: {decrypted_size}字节)")

        return final_output_path, f"解密成功！文件名: {output_filename}"

    except (InvalidToken, ValueError) as e:
        logger.error(f"解密失败: {str(e)}", exc_info=True)
        return None, f"解密失败：无效的文件格式或损坏的文件 ({str(e)})"
    except Exception as e:
        logger.error(f"解密文件时出错: {str(e)}", exc_info=True)
        return None, f"解密文件时出错: {str(e)}"
    finally:
        try:
            if os.path.exists(upload_path):
                os.remove(upload_path)
                logger.info(f"删除上传文件: {upload_path}")
        except Exception as e:
            logger.error(f"删除上传文件失败: {str(e)}")

# 自定义CSS样式
custom_css = """
body, .gradio-container {
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
    background-color: #f8f9fa;
    margin: 0;
    padding: 0;
}

/* 标题样式 */
h1 {
    text-align: center;
    color: #2c3e50;
    margin-bottom: 20px;
    font-weight: 700;
}

/* 卡片样式 */
.card {
    border-radius: 12px !important;
    box-shadow: 0 4px 20px rgba(0, 0, 0, 0.08) !important;
    border: 1px solid #eaeaea !important;
    padding: 20px;
    background: white;
    margin-bottom: 20px;
    transition: transform 0.3s ease, box-shadow 0.3s ease;
}

.card:hover {
    transform: translateY(-5px);
    box-shadow: 0 8px 25px rgba(0, 0, 0, 0.12) !important;
}

/* 标签页样式 */
.tabs {
    border-radius: 10px !important;
    overflow: hidden;
    background: #f0f2f5 !important;
    padding: 4px !important;
    margin-bottom: 20px;
}

.tab {
    border-radius: 8px !important;
    padding: 10px 20px !important;
    margin: 0 2px !important;
    transition: all 0.3s ease !important;
}

.tab.selected {
    background: linear-gradient(135deg, #6a11cb 0%, #2575fc 100%) !important;
    color: white !important;
    font-weight: 600 !important;
}

/* 按钮样式 */
button {
    background: linear-gradient(135deg, #6a11cb 0%, #2575fc 100%) !important;
    color: white !important;
    border: none !important;
    border-radius: 8px !important;
    padding: 12px 24px !important;
    font-weight: 600 !important;
    transition: all 0.3s ease !important;
    width: 100%;
    margin-top: 10px;
}

button:hover {
    transform: scale(1.05);
    box-shadow: 0 6px 15px rgba(37, 117, 252, 0.4) !important;
}

/* 文件上传区域 */
.file-upload {
    border: 2px dashed #d1d5db !important;
    border-radius: 10px !important;
    padding: 25px;
    background-color: #f9fafb !important;
    text-align: center;
    transition: all 0.3s ease;
    width: 100%;
    margin-bottom: 20px;
}

.file-upload:hover {
    border-color: #2575fc !important;
    background-color: #edf4ff !important;
}

/* 状态消息 */
.output-markdown {
    padding: 12px;
    border-radius: 8px;
    margin-top: 10px;
    text-align: center;
    font-weight: 500;
}

.success-message {
    background-color: #e6f7ee !important;
    color: #10b981 !important;
    border: 1px solid #a7f3d0 !important;
}

.error-message {
    background-color: #fef2f2 !important;
    color: #ef4444 !important;
    border: 1px solid #fecaca !important;
}

/* 使用说明 */
.instructions {
    background: linear-gradient(to right, #f8f9fa, #e9ecef) !important;
    border-left: 4px solid #2575fc !important;
    padding: 15px 20px !important;
    border-radius: 0 8px 8px 0 !important;
    margin-top: 20px;
    font-size: 14px;
    line-height: 1.6;
}

/* 文件名样式 */
.filename {
    font-weight: bold;
    color: #2575fc;
    word-break: break-all;
}

/* 响应式布局 */
@media (max-width: 768px) {
    .container {
        padding: 10px !important;
    }
    .column {
        padding: 10px !important;
    }
}
"""

# 创建 Gradio 界面
with gr.Blocks(title="安全文件加密/解密工具", css=custom_css) as demo:
    # 标题区域
    with gr.Row():
        with gr.Column():
            gr.Markdown("""
            <center>
                <h1>🔒 安全文件加密/解密工具</h1>
                <p style="color: #6b7280; max-width: 600px; margin: 0 auto;">
                    使用AES-128加密技术保护您的Markdown文件安全
                </p>
            </center>
            """)

    # 主要功能区域
    with gr.Tabs() as tabs:
        # 加密标签页
        with gr.Tab("文件加密", id="encrypt_tab"):
            with gr.Row():
                with gr.Column():
                    # 文件上传区域
                    with gr.Group(elem_classes="card"):
                        gr.Markdown("### 1️⃣ 上传Markdown文件")
                        encrypt_input = gr.File(
                            label="",
                            file_types=[".md"],
                            elem_classes="file-upload"
                        )
                        gr.Markdown("**支持的文件类型**: .md (最大4MB)")

                    # 操作按钮
                    with gr.Group(elem_classes="card"):
                        gr.Markdown("### 2️⃣ 执行加密操作")
                        encrypt_btn = gr.Button("开始加密")

                    # 示例区域
                    with gr.Group(elem_classes="card"):
                        gr.Markdown("### 🚀 快速示例")
                        example_file_path = os.path.join(BASE_DIR, "example.md")

                        if not os.path.exists(example_file_path):
                            try:
                                with open(example_file_path, "w", encoding="utf-8") as f:
                                    f.write("# 示例文件\n\n这是一个用于加密的示例Markdown文件。")
                                logger.info(f"创建示例文件: {example_file_path}")
                            except Exception as e:
                                logger.error(f"创建示例文件失败: {str(e)}")

                        gr.Examples(
                            examples=[[example_file_path]],
                            inputs=[encrypt_input],
                            outputs=[],
                            fn=lambda x: x,
                            cache_examples=False,
                            label="点击使用示例文件"
                        )

                with gr.Column():
                    # 结果区域
                    with gr.Group(elem_classes="card"):
                        gr.Markdown("### 3️⃣ 加密结果")
                        encrypt_output = gr.File(
                            label="加密后的文件",
                            interactive=False
                        )
                        encrypt_msg = gr.Markdown(
                            label="操作状态",
                            elem_classes="output-markdown"
                        )

        # 解密标签页
        with gr.Tab("文件解密", id="decrypt_tab"):
            with gr.Row():
                with gr.Column():
                    # 文件上传区域
                    with gr.Group(elem_classes="card"):
                        gr.Markdown("### 1️⃣ 上传加密文件")
                        decrypt_input = gr.File(
                            label="",
                            file_types=[".bin"],
                            elem_classes="file-upload"
                        )
                        gr.Markdown("**支持的文件类型**: .bin (最大4MB)")

                    # 操作按钮
                    with gr.Group(elem_classes="card"):
                        gr.Markdown("### 2️⃣ 执行解密操作")
                        decrypt_btn = gr.Button("开始解密")

                with gr.Column():
                    # 结果区域
                    with gr.Group(elem_classes="card"):
                        gr.Markdown("### 3️⃣ 解密结果")
                        decrypt_output = gr.File(
                            label="解密后的文件",
                            interactive=False
                        )
                        decrypt_msg = gr.Markdown(
                            label="操作状态",
                            elem_classes="output-markdown"
                        )

    # 使用说明区域
    with gr.Row():
        with gr.Column():
            with gr.Group():
                gr.Markdown("""
                <div class="instructions">
                    <h3>📖 使用说明</h3>
                    
                    1. **加密文件**:
                       - 在"文件加密"标签页上传.md文件
                       - 点击"开始加密"按钮
                       - 下载加密后的.bin文件（保留原始文件名）
                    
                    2. **解密文件**:
                       - 在"文件解密"标签页上传.bin文件
                       - 点击"开始解密"按钮
                       - 下载解密后的.md文件（恢复原始文件名）
                    
                    3. **文件名规则**:
                       - 加密后: <span class="filename">原始文件名_encrypted.bin</span>
                       - 解密后: <span class="filename">原始文件名_decrypted.md</span>
                    
                    4. **安全特性**:
                       - 所有文件在处理后60秒内自动删除
                       - 使用AES-128加密技术
                       - 每次加密使用唯一密钥
                    
                    5. **注意事项**:
                       - 最大支持4MB文件
                       - 请勿上传敏感个人信息
                </div>
                """)

    # 页脚
    with gr.Row():
        with gr.Column():
            gr.Markdown("""
            <center>
                <div style="border-top: 1px solid #eaeaea; padding-top: 15px; color: #6b7280;">
                    <p>🔐 安全文件加密工具 v1.0 | 所有文件自动销毁 | 基于Fernet加密技术</p>
                </div>
            </center>
            """)

    # 按钮事件绑定
    encrypt_btn.click(
        encrypt_file,
        inputs=[encrypt_input],
        outputs=[encrypt_output, encrypt_msg]
    )

    decrypt_btn.click(
        decrypt_file,
        inputs=[decrypt_input],
        outputs=[decrypt_output, decrypt_msg]
    )

# 获取本地IP地址
def get_local_ip():
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(("8.8.8.8", 80))
        ip = s.getsockname()[0]
        s.close()
        return ip
    except Exception:
        return "127.0.0.1"

# 启动应用
if __name__ == '__main__':
    logger.info(f"启动应用，操作系统: {platform.system()} {platform.release()}")
    logger.info(f"Python版本: {platform.python_version()}")
    logger.info(f"Gradio版本: {gr.__version__}")

    local_ip = get_local_ip()
    local_url = f"http://localhost:5001"
    network_url = f"http://{local_ip}:5001"

    print("\n" + "=" * 60)
    print(f"应用已启动！请访问以下URL:")
    print(f"本地访问: {local_url}")
    print(f"局域网访问: {network_url}")
    print("=" * 60 + "\n")

    try:
        # 绑定到 0.0.0.0 地址，允许局域网访问
        demo.launch(
            server_port=5001,
            server_name="0.0.0.0",  # 保持这个配置不变
            show_error=True
        )
    except Exception as e:
        logger.error(f"启动应用失败: {str(e)}", exc_info=True)
        raise